#pragma once
#include "yzrutil.h"
#include <cmath>
#include <vector>
#include <memory>
#include "lang/Math.h"

namespace yzrilyzr_dsp{
	EBCLASS(DCT){
	private:
	int32_t N;
	float alpha, alpha0;
	std::vector<std::vector<float>> cosTable; // 预计算的余弦表
	std::vector<float> cuTable;   // 预计算的cu系数表
	std::vector<float> cvTable;   // 预计算的cv系数表

	public:
	DCT(int32_t n) : N(n){
		if(N <= 0) throw std::invalid_argument("N must be positive");

		alpha0=1.0f / std::sqrt(2.0f);
		alpha=std::sqrt(2.0f / N);

		// 预计算余弦表和系数表
		cosTable.resize(N, std::vector<float>(N * 2));
		cuTable.resize(N);
		cvTable.resize(N);

		for(u_index u=0; u < N; ++u){
			cuTable[u]=(u == 0)?alpha0:alpha;
			for(u_index x=0; x < N; ++x){
				cosTable[u][x]=std::cos((2 * x + 1) * u * yzrilyzr_lang::Math::PI / (2 * N));
			}
		}

		for(u_index v=0; v < N; ++v){
			cvTable[v]=(v == 0)?alpha0:alpha;
		}
	}

	// 使用一维数组接口提高缓存利用率
	void dct(const float * input, float * output){
		std::vector<float> temp(N * N);

		// 先做行变换
		for(u_index y=0; y < N; ++y){
			for(u_index u=0; u < N; ++u){
				float sum=0.0f;
				for(u_index x=0; x < N; ++x){
					sum+=input[y * N + x] * cosTable[u][x];
				}
				temp[y * N + u]=sum * cuTable[u];
			}
		}

		// 再做列变换
		for(u_index u=0; u < N; ++u){
			for(u_index v=0; v < N; ++v){
				float sum=0.0f;
				for(u_index y=0; y < N; ++y){
					sum+=temp[y * N + u] * cosTable[v][y];
				}
				output[u * N + v]=sum * cvTable[v];
			}
		}
	}

	void idct(const float * input, float * output){
		std::vector<float> temp(N * N);

		// 先做行变换
		for(u_index v=0; v < N; ++v){
			for(u_index y=0; y < N; ++y){
				float sum=0.0f;
				for(u_index u=0; u < N; ++u){
					sum+=input[v * N + u] * cuTable[u] * cosTable[u][y];
				}
				temp[v * N + y]=sum;
			}
		}

		// 再做列变换
		for(u_index y=0; y < N; ++y){
			for(u_index x=0; x < N; ++x){
				float sum=0.0f;
				for(u_index v=0; v < N; ++v){
					sum+=temp[v * N + y] * cvTable[v] * cosTable[v][x];
				}
				output[y * N + x]=sum;
			}
		}
	}
	};
}